#pragma once

#include "llvm/Module.h"
#include "llvm/Function.h"
#include "llvm/BasicBlock.h"
#include "llvm/Instruction.h"
#include "llvm/Type.h"
#include "llvm/DerivedTypes.h"
#include "llvm/ExecutionEngine/JIT.h"
#include "llvm/Instructions.h"
#include "llvm/Support/raw_ostream.h"
#include "Utils.h"
#include <vector>
#include <map>
#include "Callable.h"

using namespace llvm;
using namespace std;

class InstrumenterBase
{
public:
	ExecutionEngine *EE;
protected:
	Module *M;
	LLVMContext &C;
	Value *Dummy;

	InstrumenterBase(Module *M, ExecutionEngine *EE) :
		EE(EE),
		M(M),
		C(M->getContext()),
		Dummy(ConstantInt::getNullValue(Type::getInt32Ty(C)))
	{
	}
};

template<typename InstructionType, typename ProfilerType, typename ArgType=int>
class Instrumenter : public InstrumenterBase
{
public:
	template<typename Collection>
	void Instrument(Collection &Instructions)
	{
		Func.AddType(getArgType()).AddType(Type::getInt8PtrTy(C)).Map(mighty_cast<void*>(&Callback), EE, "");

		const Type *PVoid = Type::getInt8PtrTy(C), *PWord = Type::getInt64Ty(C);

		for (typename Collection::iterator it = Instructions.begin(), ite = Instructions.end(); it != ite; ++it)
			if (InstructionType *t = dyn_cast<InstructionType>(*it))
				if (Value *V = InstrumentThis(t)) {
					ProfilerType *Profiler = new ProfilerType(t, this);
					Value *Args[] = {V, ConstantExpr::getIntToPtr(ConstantInt::get(PWord, (int64_t)Profiler), PVoid)};
					Func.Call(Args, t);
				}
	}
protected:
	virtual Value* InstrumentThis(InstructionType *Inst) = 0;
	virtual const Type* getArgType() { return Type::getInt32Ty(C); }

	Instrumenter(Module *M, ExecutionEngine *EE) : InstrumenterBase(M, EE), Func(M, Type::getVoidTy(C)) {}
private:
	Callable Func;

	static void Callback(ArgType Arg, void *PPDT)
	{
		static_cast<ProfilerType*>(PPDT)->Profile(Arg);
	}
};
